import os
import pickle

from PIL import Image
from tqdm import tqdm

from config.config import DATA_PROCESSED
from data.utils import get_class_names


def process_batch(batch_path, output_dir, class_names):
    """处理单个批次数据"""
    with open(batch_path, 'rb') as f:
        batch = pickle.load(f, encoding='bytes')

    images = batch[b'data']
    labels = batch[b'labels']
    filenames = batch[b'filenames']

    for idx in tqdm(range(len(images)), desc=f"Processing {batch_path.name}"):
        # 转换图像格式 (CIFAR-10的特殊存储格式)
        img_flat = images[idx]
        img = img_flat.reshape(3, 32, 32).transpose(1, 2, 0)
        img = Image.fromarray(img)

        # 获取类别信息
        label_idx = labels[idx]
        class_name = class_names[label_idx]
        filename = filenames[idx].decode('utf-8')

        # 创建类别目录并保存图像
        class_dir = output_dir / class_name
        class_dir.mkdir(parents=True, exist_ok=True)
        img.save(class_dir / filename)


def extract_images(raw_dir):
    """主提取函数"""
    # 获取类别名称
    class_names = get_class_names(raw_dir)

    # 创建输出目录
    train_dir = DATA_PROCESSED / "train"
    test_dir = DATA_PROCESSED / "test"

    if not os.path.exists(train_dir):
        # 处理训练数据 (5个批次)
        print("Processing training data...")
        for i in range(1, 6):
            batch_path = raw_dir / f"data_batch_{i}"
            process_batch(batch_path, train_dir, class_names)
    else:
        print("Train files already extracted. Skipping processing.")

    if not os.path.exists(test_dir):
        # 处理测试数据
        print("Processing test data...")
        test_batch = raw_dir / "test_batch"
        process_batch(test_batch, test_dir, class_names)

        print(f"Data extracted to:\nTrain: {train_dir}\nTest: {test_dir}")
    else:
        print("Test files already extracted. Skipping processing.")
